import pandas as pd
import os
from common.mongo_utils import get_db
from engine.data.expr.expr_mgr import ExprMgr
from datetime import datetime
from common import config

class Dataloader:
    def __init__(self, path=None):
        curr_date = datetime.now().strftime('%Y%m%d')
        if not path:
            path = config.get_config_path() + '/csv'
        if not os.path.exists(path):
            os.mkdir(path)

        self.path = path
        self.dict_code_df = {}
        self.expr = ExprMgr()
        self.expr.init()

    def _clean_df(self, df):
        if len(df) == 0:
            print('__clean_df: df is none')
            return
        df.rename(columns={'trade_date': 'date', 'ts_code': 'code', 'vol': 'volume','nav_date':'date'}, inplace=True)
        df['date'] = df['date'].apply(lambda x: str(x))



        #df = df[['open', 'high', 'low', 'close', 'date', 'code', 'volume', 'amount','adj_factor']]
        if 'adj_factor' in df.columns:
            for col in ['open','high','low','close']:
                df[col] = df[col] * df['adj_factor']
        df.index = df['date']
        df.sort_index(ascending=True, inplace=True)
        if 'close' in df.columns:
            df['rate'] = df['close'].pct_change()
        if 'adj_nav' in df.columns:
            df['rate'] = df['adj_nav'].pct_change()
        #df['rate'] = df['close'].pct_change()
        return df

    def _load_from_file(self, code):
        filename = '{}/{}.csv'.format(self.path, code)
        if not os.path.exists(filename):
            print('csv不存在！')
            #return None
            df = self._load_from_mongo(code)
        else:
            df = pd.read_csv(filename, index_col=[0])  # index_col指定是第一列，这是保存的时候生成的。
        df = self._clean_df(df)
        return df

    def _load_from_mongo(self, code):
        filter = {'_id': 0}
        for col in ['open', 'high', 'low', 'close', 'vol', 'ts_code', 'amount', 'trade_date', 'adj_factor','adj_nav','nav_date']:
            filter[col] = 1

        tb = 'quotes'
        if '.' not in code or code in config.get_index_codes():
            tb = 'quotes_index'
        if '.OF' in code:
            tb = 'navs'

        items = get_db()[tb].find({'ts_code': code}, filter)
        df = pd.DataFrame(list(items))
        df.to_csv(self.path + '/{}.csv'.format(code))
        return df

    def load_code_df(self, code):
        if code in self.dict_code_df.keys():
            return self.dict_code_df[code]
        else:
            df = self._load_from_file(code)
            self.dict_code_df[code] = df
        return df

    def calc_code_features(self, code ,fields, names):
        df = self.load_code_df(code)
        for feature, name in zip(fields, names):
            if feature == "":
                continue

            if name in df.columns:
                continue
            else:
                expr = self.expr.get_expression(feature)
                se = expr.load(code)
                se.name = name
                df[name] = se
        self.dict_code_df[code] = df

    def get_codes(self, codes):
        dfs = []
        for code in codes:
            df = self.load_code_df(code)
            df.dropna(inplace=True)
            dfs.append(df)
        return dfs

    def load(self, codes, start_time=None, end_time=datetime.now().strftime('%Y%m%d'), fields=[], names=None):
        if len(fields) > 0:
            if fields and not names:
                names = fields
            for code in codes:
                # 这里未加载即加载过程
                print('加载',code)
                self.calc_code_features(code, fields, names)

        dfs = self.get_codes(codes)

        if not start_time:
            start_time = self._calc_max_date(dfs)

        dfs = [df.loc[start_time:end_time] for df in dfs]
        start_time_new = self._calc_max_date(dfs)
        if start_time_new > start_time:
            dfs = [df.loc[start_time_new:end_time] for df in dfs]

        calendar = self._calc_calendar(dfs)
        calendar.sort()
        #print(calendar, len(calendar))
        dfs = self._reindex_dfs(calendar, dfs) #返回一个list

        df_all = pd.concat(dfs, axis=0)
        df_all.sort_index(inplace=True)

        print(start_time, end_time)
        df_all = df_all.loc[start_time:end_time]
        return df_all

    def _reindex_dfs(self, calendar, dfs):
        new_dfs = []
        for df in dfs:
            old_index = df.index
            diff = list(set(calendar).difference(set(old_index)))
            #print('差值：',diff)
            new_df = df.copy(deep=True)
            del new_df['date']
            new_df = new_df.reindex(index=calendar, method='ffill')
            new_df['date'] = new_df.index
            #new_df['rate'] = new_df['close'].pct_change()
            new_dfs.append(new_df)
        return new_dfs

    def _calc_calendar(self, dfs):
        calenlar = None
        for df in dfs:
            indexes = list(df.index)
            if calenlar is None:
                calenlar = indexes
            else:
                calenlar = list(set(calenlar).union(set(indexes)))
        return calenlar

    def _calc_max_date(self, dfs):
        max_date = '19900101'
        for df in dfs:
            date = df.index[0]
            if date > max_date:
                max_date = date
        return max_date

D = Dataloader()

if __name__ == '__main__':
    fields = ['Return($close,20)','Ref($close,-5)/$close -1']
    names = ['return_20','label']
    df = D.load(['000300.SH'], fields=fields, names=names)
    print(df.head(20))
    #df.dropna(inplace=True)
    #print(df[['rate','close','date','code']])